#coding=utf-8
import cv2
import numpy as np
import mvsdk
import time
import platform
from ctypes import c_uint, byref

class CameraController(object):
	def __init__(self):
		super(CameraController, self).__init__()
		self.pFrameBuffer = 0
		self.quit = False
		self.hCamera = 0
		self.cap = None
		self.monoCamera = False

	def enumerate_cameras(self):
		"""枚举相机设备"""
		DevList = mvsdk.CameraEnumerateDevice()
		nDev = len(DevList)
		if nDev < 1:
			print("未找到相机设备!")
			return None
		
		print("找到 {} 个相机设备:".format(nDev))
		for i, DevInfo in enumerate(DevList):
			print("{}: {} {}".format(i, DevInfo.GetFriendlyName(), DevInfo.GetPortType()))
		
		return DevList

	def open_camera(self, dev_index=0):
		"""打开指定索引的相机"""
		DevList = self.enumerate_cameras()
		if not DevList:
			return False
		
		if dev_index >= len(DevList):
			print("设备索引超出范围!")
			return False
		
		DevInfo = DevList[dev_index]
		print("正在打开相机: {}".format(DevInfo.GetFriendlyName()))
		
		try:
			self.hCamera = mvsdk.CameraInit(DevInfo, -1, -1)
		except mvsdk.CameraException as e:
			print("相机初始化失败({}): {}".format(e.error_code, e.message))
			return False
		
		# 获取相机特性描述
		self.cap = mvsdk.CameraGetCapability(self.hCamera)
		
		# 判断是黑白相机还是彩色相机
		self.monoCamera = (self.cap.sIspCapacity.bMonoSensor != 0)
		
		return True

	def setup_camera_parameters(self):
		"""设置相机基本参数"""
		if not self.hCamera:
			print("相机未打开!")
			return False
		
		# 设置图像输出格式
		if self.monoCamera:
			mvsdk.CameraSetIspOutFormat(self.hCamera, mvsdk.CAMERA_MEDIA_TYPE_MONO8)
		else:
			mvsdk.CameraSetIspOutFormat(self.hCamera, mvsdk.CAMERA_MEDIA_TYPE_BGR8)
		
		# 设置连续采集模式
		mvsdk.CameraSetTriggerMode(self.hCamera, 0)
		
		# 手动曝光设置
		mvsdk.CameraSetAeState(self.hCamera, 0)
		mvsdk.CameraSetExposureTime(self.hCamera, 30 * 1000)
		
		return True

	def set_roi(self, x, y, width, height):
		"""设置ROI区域"""
		if not self.hCamera:
			print("相机未打开!")
			return False
		
		# 获取最大分辨率
		max_width = self.cap.sResolutionRange.iWidthMax
		max_height = self.cap.sResolutionRange.iHeightMax
		
		# 检查ROI参数是否有效
		if x < 0 or y < 0 or width <= 0 or height <= 0:
			print("ROI参数无效!")
			return False
		
		if x + width > max_width or y + height > max_height:
			print("ROI超出相机最大分辨率范围!")
			print("最大分辨率: {}x{}".format(max_width, max_height))
			return False
		
		try:
			# 停止当前采集
			mvsdk.CameraPause(self.hCamera)
			
			# 使用自定义分辨率设置ROI区域
			# 0xFF表示自定义分辨率(ROI)
			resolution = mvsdk.tSdkImageResolution()
			resolution.iIndex = 0xFF  # 自定义分辨率
			resolution.iHOffsetFOV = x
			resolution.iVOffsetFOV = y
			resolution.iWidthFOV = width
			resolution.iHeightFOV = height
			resolution.iWidth = width
			resolution.iHeight = height
			
			mvsdk.CameraSetImageResolution(self.hCamera, resolution)
			
			# 重新分配帧缓冲区以适应新的ROI尺寸
			if self.pFrameBuffer:
				mvsdk.CameraAlignFree(self.pFrameBuffer)
				self.pFrameBuffer = 0
			
			# 重新分配帧缓冲区
			FrameBufferSize = width * height * (1 if self.monoCamera else 3)
			self.pFrameBuffer = mvsdk.CameraAlignMalloc(FrameBufferSize, 16)
			
			# 重新开始采集
			mvsdk.CameraPlay(self.hCamera)
			
			print("ROI设置成功: x={}, y={}, width={}, height={}".format(x, y, width, height))
			return True
		except mvsdk.CameraException as e:
			print("ROI设置失败({}): {}".format(e.error_code, e.message))
			# 如果设置失败，尝试恢复采集
			try:
				mvsdk.CameraPlay(self.hCamera)
			except:
				pass
			return False

	def get_current_roi(self):
		"""获取当前ROI设置"""
		if not self.hCamera:
			print("相机未打开!")
			return None
		
		try:
			# 获取当前分辨率设置
			resolution = mvsdk.CameraGetImageResolution(self.hCamera)
			
			roi_info = {
				'x': resolution.iHOffsetFOV,
				'y': resolution.iVOffsetFOV,
				'width': resolution.iWidth,
				'height': resolution.iHeight
			}
			
			return roi_info
		except mvsdk.CameraException as e:
			print("获取ROI信息失败({}): {}".format(e.error_code, e.message))
			return None

	def allocate_frame_buffer(self):
		"""分配帧缓冲区"""
		if not self.cap:
			print("相机能力信息未获取!")
			return False
		
		# 计算RGB buffer所需的大小
		FrameBufferSize = self.cap.sResolutionRange.iWidthMax * self.cap.sResolutionRange.iHeightMax * (1 if self.monoCamera else 3)
		
		# 分配RGB buffer
		self.pFrameBuffer = mvsdk.CameraAlignMalloc(FrameBufferSize, 16)
		
		return True

	def capture_image(self, filename=None):
		"""捕获一张图像并保存"""
		if not self.hCamera:
			print("相机未打开!")
			return False
		
		try:
			# 分配临时帧缓冲区
			FrameBufferSize = self.cap.sResolutionRange.iWidthMax * self.cap.sResolutionRange.iHeightMax * (1 if self.monoCamera else 3)
			pFrameBuffer = mvsdk.CameraAlignMalloc(FrameBufferSize, 16)
			
			# 抓取一帧图像
			pRawData, FrameHead = mvsdk.CameraGetImageBuffer(self.hCamera, 2000)
			
			# 图像处理
			mvsdk.CameraImageProcess(self.hCamera, pRawData, pFrameBuffer, FrameHead)
			mvsdk.CameraReleaseImageBuffer(self.hCamera, pRawData)

			# Windows平台图像翻转
			if platform.system() == "Windows":
				mvsdk.CameraFlipFrameBuffer(pFrameBuffer, FrameHead, 1)
			
			# 转换为OpenCV格式
			frame_data = (mvsdk.c_ubyte * FrameHead.uBytes).from_address(pFrameBuffer)
			frame = np.frombuffer(frame_data, dtype=np.uint8)
			frame = frame.reshape((FrameHead.iHeight, FrameHead.iWidth, 1 if FrameHead.uiMediaType == mvsdk.CAMERA_MEDIA_TYPE_MONO8 else 3))
			
			# 生成文件名
			if filename is None:
				import datetime
				timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
				roi_info = self.get_current_roi()
				if roi_info:
					filename = "capture_{}_{}x{}+{}+{}.png".format(
						timestamp, roi_info['width'], roi_info['height'], roi_info['x'], roi_info['y'])
				else:
					filename = "capture_{}.png".format(timestamp)
			
			# 保存图像
			cv2.imwrite(filename, frame)
			print("图像已保存: {}".format(filename))
			print("图像尺寸: {}x{}".format(FrameHead.iWidth, FrameHead.iHeight))
			
			# 释放临时缓冲区
			mvsdk.CameraAlignFree(pFrameBuffer)
			
			return True
			
		except mvsdk.CameraException as e:
			print("图像捕获失败({}): {}".format(e.error_code, e.message))
			return False

	def capture_after_roi_setup(self):
		"""设置ROI后拍摄图像"""
		if not self.hCamera:
			print("相机未打开!")
			return
		
		# 启动相机
		mvsdk.CameraPlay(self.hCamera)
		
		print("相机已启动，准备拍摄...")
		
		# 等待相机稳定
		time.sleep(0.5)
		
		# 拍摄图像
		success = self.capture_image()
		
		if success:
			print("拍摄完成!")
		else:
			print("拍摄失败!")

	def close_camera(self):
		"""关闭相机"""
		if self.hCamera:
			mvsdk.CameraUnInit(self.hCamera)
			self.hCamera = 0
		
		if self.pFrameBuffer:
			mvsdk.CameraAlignFree(self.pFrameBuffer)
			self.pFrameBuffer = 0

	def interactive_roi_setup(self):
		"""交互式ROI设置"""
		if not self.hCamera:
			print("相机未打开!")
			return
		
		# 获取当前ROI信息
		current_roi = self.get_current_roi()
		if current_roi:
			print("当前ROI设置: x={}, y={}, width={}, height={}".format(
				current_roi['x'], current_roi['y'], current_roi['width'], current_roi['height']))
		
		# 获取最大分辨率
		max_width = self.cap.sResolutionRange.iWidthMax
		max_height = self.cap.sResolutionRange.iHeightMax
		print("相机最大分辨率: {}x{}".format(max_width, max_height))
		
		# 交互式设置ROI
		print("\n设置ROI参数:")
		try:
			x = int(input("请输入ROI左上角X坐标 (默认0): ") or 0)
			y = int(input("请输入ROI左上角Y坐标 (默认0): ") or 0)
			width = int(input("请输入ROI宽度 (默认{}): ".format(max_width)) or max_width)
			height = int(input("请输入ROI高度 (默认{}): ".format(max_height)) or max_height)
			
			if self.set_roi(x, y, width, height):
				print("ROI设置完成!")
			else:
				print("ROI设置失败!")
			
		except ValueError:
			print("输入参数无效!")


def main():
	controller = CameraController()
	
	# 打开相机
	if not controller.open_camera():
		return
	
	try:
		# 设置相机基本参数
		if not controller.setup_camera_parameters():
			return
		
		# 交互式ROI设置
		controller.interactive_roi_setup()
		
		# 设置ROI后拍摄图像
		controller.capture_after_roi_setup()
		
	finally:
		# 关闭相机
		controller.close_camera()

if __name__ == "__main__":
	main()